/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.hylanda.processors.imagerecognition;

import org.apache.nifi.components.PropertyDescriptor;
import org.apache.nifi.flowfile.FlowFile;
import org.apache.nifi.flowfile.attributes.CoreAttributes;
import org.apache.commons.io.IOUtils;
import org.apache.nifi.annotation.behavior.ReadsAttribute;
import org.apache.nifi.annotation.behavior.ReadsAttributes;
import org.apache.nifi.annotation.behavior.WritesAttribute;
import org.apache.nifi.annotation.behavior.WritesAttributes;
import org.apache.nifi.annotation.lifecycle.OnScheduled;
import org.apache.nifi.annotation.documentation.CapabilityDescription;
import org.apache.nifi.annotation.documentation.SeeAlso;
import org.apache.nifi.annotation.documentation.Tags;
import org.apache.nifi.processor.exception.ProcessException;
import org.apache.nifi.processor.AbstractProcessor;
import org.apache.nifi.processor.ProcessContext;
import org.apache.nifi.processor.ProcessSession;
import org.apache.nifi.processor.ProcessorInitializationContext;
import org.apache.nifi.processor.Relationship;
import org.apache.nifi.processor.util.StandardValidators;

import com.alibaba.fastjson.JSONObject;
import com.hylanda.processors.imagerecognition.util.GenUrlCrc64;
import com.hylanda.processors.imagerecognition.util.HTTPPost;
import com.hylanda.processors.imagerecognition.util.NifiUtil;

import java.io.BufferedOutputStream;
import java.io.OutputStream;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.UUID;

@Tags({ "image text recognition" })
@CapabilityDescription("文字识别")
@SeeAlso({})
@ReadsAttributes({ @ReadsAttribute(attribute = "", description = "") })
@WritesAttributes({ @WritesAttribute(attribute = "", description = "") })
public class GetImageTextProcessor extends AbstractProcessor {

	public static final PropertyDescriptor  API_KEY = new PropertyDescriptor.Builder().name("API_KEY")
			.displayName("API_KEY").description("百度账号下图片识别应用的apikey,注册地址：https://cloud.baidu.com/product/ocr").addValidator(StandardValidators.NON_EMPTY_VALIDATOR).build();
	public static final PropertyDescriptor  SECRET_KEY = new PropertyDescriptor.Builder().name("SECRET_KEY")
			.displayName("SECRET_KEY").description("百度账号下图片识别应用的secrey key,注册地址：https://cloud.baidu.com/product/ocr").addValidator(StandardValidators.NON_EMPTY_VALIDATOR).build();
	/*public static final PropertyDescriptor IMAGE_URL = new PropertyDescriptor.Builder().name("IMAGE_URL")
			.displayName("IMAGE_URL").description("image url").required(false)
			.build();*/
	public static final PropertyDescriptor IMAGE_FIELD = new PropertyDescriptor.Builder().name("IMAGE_FIELD")
			.displayName("图片url所在的字段").description("image filed,默认是image_url").addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
			.build();
	public static final PropertyDescriptor MODE = new PropertyDescriptor
    		.Builder().name("MODE")
    		.displayName("识别模式")
    		.description("图片识别模式")
    		.required(true)
    		.allowableValues(new String[] { "通用文字识别","高精准","含位置信息","含位置高精度","含生僻字","网络图片文字识别" })
    		.addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
    		.build();

	public static final Relationship SUCCESS = new Relationship.Builder().name("SUCCESS").description("SUCCESS")
			.build();
	public static final Relationship FAILURE = new Relationship.Builder().name("FAILURE").description("FAILURE")
			.build();

	private List<PropertyDescriptor> descriptors;

	private Set<Relationship> relationships;

	@Override
	protected void init(final ProcessorInitializationContext context) {
		final List<PropertyDescriptor> descriptors = new ArrayList<PropertyDescriptor>();
		
		//descriptors.add(IMAGE_URL);
		descriptors.add(MODE);
		descriptors.add(API_KEY);
		descriptors.add(SECRET_KEY);
		descriptors.add(IMAGE_FIELD);
		this.descriptors = Collections.unmodifiableList(descriptors);

		final Set<Relationship> relationships = new HashSet<Relationship>();
		relationships.add(SUCCESS);
		relationships.add(FAILURE);
		this.relationships = Collections.unmodifiableSet(relationships);
	}

	@Override
	public Set<Relationship> getRelationships() {
		return this.relationships;
	}

	@Override
	public final List<PropertyDescriptor> getSupportedPropertyDescriptors() {
		return descriptors;
	}

	@OnScheduled
	public void onScheduled(final ProcessContext context) {

	}

	@Override
	public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
		FlowFile flowFile = session.get();
		if (flowFile == null) {
			return;
		}
		String content = NifiUtil.getContent(session, flowFile);
		String mode = null;
		try {
			mode = context.getProperty(MODE).getValue().trim();
		} catch (Exception e) {
			// TODO: handle exception
		}
		String ak = null;
		try {
			ak = context.getProperty(API_KEY).getValue().trim();
		} catch (Exception e) {
			// TODO: handle exception
		}
		String sk = null;
		try {
			sk = context.getProperty(SECRET_KEY).getValue().trim();
		} catch (Exception e) {
			// TODO: handle exception
		}
		
		String imageFiled = null;
		try {
			imageFiled = context.getProperty(IMAGE_FIELD).getValue().trim();
		} catch (Exception e) {
			// TODO: handle exception
		}
		if (null == imageFiled||imageFiled.equals("")) {
			imageFiled = "image_url";
		}
		String imageUrl = null;
		JSONObject sourceJSON = JSONObject.parseObject(content);
		imageUrl = sourceJSON.getString(imageFiled);
		
		if (null==imageUrl||imageUrl.equals("")) {
        	flowFile = session.write(flowFile, (in, out) -> {
				try (OutputStream outputStream = new BufferedOutputStream(out)) {				
					outputStream.write(IOUtils.toByteArray("图片url为空"));
				}
			});
        	session.transfer(flowFile, FAILURE);
			return;
		}
		if (null==ak||"".equals(ak)||null==sk||"".equals(sk)) {
			//String clientId = "4MLCVPbPlrggrWXsECVA0K4j";
	        // 官网获取的 Secret Key 更新为你注册的
	        //String clientSecret = "dYdnIwGr21oRKjl768d3oelcDfNL9osm";
			ak = "4MLCVPbPlrggrWXsECVA0K4j";
			sk = "dYdnIwGr21oRKjl768d3oelcDfNL9osm";
		}
		int modeint = 1;
		switch (mode) {//"高精准","含位置信息","含位置高精度","含生僻字","网络图片文字识别"
		case "通用文字识别":
			modeint = 1;
			break;
		case "高精准":
			modeint = 2;
			break;
		case "含位置信息":
			modeint = 3;
			break;
		case "含位置高精度":
			modeint = 4;
			break;
		case "含生僻字":
			modeint = 5;
			break;
		case "网络图片文字识别":
			modeint = 6;
			break;

		default:
			break;
		}
		
		String result = HTTPPost.getTextFromImage(imageUrl, modeint,ak,sk);
		flowFile = session.write(flowFile, (in, out) -> {
			try (OutputStream outputStream = new BufferedOutputStream(out)) {
				outputStream.write(IOUtils.toByteArray(result));
			}
		});

		String uuid = GenUrlCrc64.GenCrc64(UUID.randomUUID().toString());
		String session_id = flowFile.getAttribute(CoreAttributes.UUID.key());
		String instance_id = getIdentifier();
		JSONObject logJSON = new JSONObject();
		JSONObject msgJSON = new JSONObject();

		Date date = new Date();
		SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
		String time = sdf.format(date);
		msgJSON.put("time", time);
		msgJSON.put("status", "success");

		logJSON.put("msg", msgJSON);
		logJSON.put("session_id", session_id);
		logJSON.put("url", imageUrl);
		logJSON.put("uuid", uuid);
		logJSON.put("instance_id", instance_id);

		getLogger().info(logJSON.toJSONString());
		session.transfer(flowFile, SUCCESS);
	}
}
